Add force_fused option to scaled_dot_product_attention - #4185

Proposed changes

Add a force_fused option to mx.fast.scaled_dot_product_attention that
ignores the automatic dispatch heuristics, uses a fused kernel, and raises
with the unsupported constraint when no fused kernel is available. This lets
inference runtimes — the only party that knows their resident weights, KV
cache size, and memory budget — decide when the fused path's bounded-memory
behavior is worth its throughput cost.

Also restores the fused full-attention support for head_dim=192/256 that was
proposed in #3293/#3660. Those kernels are now reachable only through
force_fused; the default dispatch does not route to them.

Follow-up hardening incorporates the full-dispatch threadgroup safety check
and capability-diagnostic concerns raised by @apocryphx in #4186, while
retaining this PR's 192/256 kernels and C++ positional-call compatibility.

Background

#3658 tracks models with genuine head_dim=192/256 full-attention layers. At
long context, the unfused path materializes a score
transient of O(n_heads × qL × kL) per full-attention layer. The fused
full-attention kernel tiles K/V so its transient stays bounded regardless of
kL. However, benchmark measurements on a genuine head_dim=256 workload
(the Qwen 3.6 35B checkpoint measured in #3658) show the fused path trades
prompt throughput for that memory floor:

These stock mlx-lm benchmark arms were separated by other main-branch
commits, so they are supporting workload evidence rather than a strict
single-commit A/B. The optional 131K fused arm aborted with an
interactivity-impacting command-buffer error and is intentionally omitted.

A fixed dispatch threshold cannot pick correctly for all runtimes, so the
decision is exposed to the caller instead. This PR is the follow-up to the
maintainer request in #3658:

Let's add a force_fused=True option to mx.fast.scaled_dot_product_attention API that ignores the heuristics and uses fused path, and throws when a fused kernel is not available for the shape. — @zcbenz

Supersedes the closed #3660.
Addresses #3658.

API

mx.fast.scaled_dot_product_attention(
    q, k, v, *, scale, mask=None, sinks=None, force_fused=False, stream=None
)

• force_fused=False (default): preserves existing dispatch for previously
  supported shapes. The missing head_dim=192 vector kernel is added but
  reachable only through force_fused.
• force_fused=True: bypass the dispatch heuristics and use a fused kernel.
  Raises ValueError with the unsupported constraint when no fused kernel is
  available, which covers:
  • shapes not covered by the full- or vector-attention fused kernels
    (including unsupported head_dims, e.g. 512);
  • CPU execution;
  • training mode that requires a logsumexp output (the fused kernels do not
    emit one).

is_training on its own (when the VJP path already falls back) is allowed to
use the fused forward kernel — the output is identical.

Kernel support

• Full-attention (steel_attention): add head_dim=192/256 instantiations.
  Selected only when force_fused=True. Extend the existing large-head-dim
  V-tile synchronization barriers from BD == 128 to BD >= 128, covering
  192/256 as well as 128. This synchronization gap was also independently
  identified in giaki3003/mlx@1175b4d.
• Vector (decode) (sdpa_vector): add the missing head_dim=192 (192,192)
  and aggregation instantiations. The new 192 vector kernel is opt-in via
  force_fused like the new full-attention kernels; the existing head_dim=256
  vector kernel keeps its automatic routing.
• Route head_dim >= 192 away from the NAX kernel family, which has no
  192/256 instantiations.
• Check the compiled pipeline's maximum threadgroup size before both full
  attention dispatches, avoiding silent zero output on register-limited GPUs.
• Use each backend's actual support matrix: Metal full/vector kernels and CUDA
  cuDNN/vector kernels share the same opt-in semantics.

Validation

• Focused tests in python/tests/test_fast_sdpa.py:
  • force_fused=True full attention for head_dim 192/256 at short key
    lengths (where the default dispatch stays unfused) matches the reference;
  • force_fused=True full attention for head_dim 192/256 at kL=16385
    matches the reference, covering the Steel V-tile barrier path;
  • force_fused=True on an already-supported head_dim=128 matches;
  • force_fused=True on an unsupported shape (head_dim=512) raises;
  • unsupported full/vector head dims, causal qL > kL, vector GQA limits,
    and CPU execution raise with the specific violated constraint;
  • default dispatch with head_dim 192/256 full attention stays unfused and
    matches the reference;
  • head_dim=192 vector attention matches the reference.
• Existing test_fast_sdpa.py and test_fast.py suites continue to pass.

Scope

• Gemma 4 26B/31B are not affected: their long-context global layers use
  head_dim=512, which is outside this PR's scope (see Metal SDPA: no fused path for head_dim=512 (vector or full kernel), and the fallback is silent #3885).
  head_dim=256 layers there are bounded sliding-window layers.
• The Qwen 3.6 35B checkpoint measured in Allow force using fused kernel in scaled_dot_product_attention #3658 uses genuine
  head_dim=256 full attention and demonstrates the intended memory/throughput
  tradeoff. The optional 131K fused arm did not complete, so this PR does not
  claim a valid 131K A/B result.
• CUDA uses its cuDNN/vector support matrix for force_fused; a no-GPU build
  raises rather than silently ignoring the flag. CUDA compilation remains a
  GitHub CI validation boundary because the local validation host is macOS/Metal.

Checklist

• I have read the CONTRIBUTING document
• I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
• I have added tests that prove my fix is effective or that my feature works
• I have updated the necessary documentation (API docstring)